/*
 * Codegen.cpp
 *
 *  Created on: 1 Nov 2010
 *      Author: kn308
 */

#include <sstream>
#include <iomanip>

#include <MAlice/MAlice.h>
#include <MAlice/Codegen.h>
#include <MAlice/MAliceDriver.hpp>
#include <MAlice/Ast/StringLiteralAst.h>
#include <MAlice/Ast/IntegerLiteralAst.h>
#include <MAlice/Ast/CharacterLiteralAst.h>

#include <llvm/LLVMContext.h>
#include <llvm/Module.h>
#include <llvm/PassManager.h>
#include <llvm/Instructions.h>
#include <llvm/Pass.h>
#include <llvm/Function.h>

#include <llvm/Support/raw_ostream.h>

#include <llvm/Analysis/Verifier.h>

#include <llvm/CodeGen/DwarfWriter.h>

#include <llvm/Bitcode/Archive.h>
#include <llvm/Bitcode/ReaderWriter.h>
#include <llvm/Bitcode/BitstreamReader.h>
#include <llvm/Bitcode/BitstreamWriter.h>

#include <llvm/Target/TargetData.h>
#include <llvm/Target/TargetSelect.h>

#include <llvm/Transforms/Scalar.h>
#include <llvm/Target/TargetSelect.h>
#include <llvm/TypeSymbolTable.h>

#include <llvm/Support/IRBuilder.h>


#include <MAlice/Ast/TypeAst.h>
#include <MAlice/Ast/PointerTypeAst.h>

#include <MAlice/Ast/ListAst.h>
#include <MAlice/Ast/FunctionProtoAst.h>
#include <MAlice/Ast/FunctionParameterAst.h>
#include <MAlice/Ast/FunctionDeclAst.h>

static MAlice::Codegen* gCodegen = NULL;
namespace MAlice
{

/* auto reverse class */
AutoReverse::AutoReverse(Codegen* g) :
	mG(g)
{
	mBB = mG->getBuilder()->GetInsertBlock();
	mG->pushScope();
}

AutoReverse::~AutoReverse()
{
	mG->popScope();

	// if there is not basic block prior to this auto reverse, there is nothing to be popped!
	if(!mBB) return;
	mG->getBuilder()->SetInsertPoint(mBB);
}


static const int gEntryBlockFinalizerPid = 0;
class EntryBlockFinalizer : public llvm::FunctionPass
{
public:
	EntryBlockFinalizer()
	 : FunctionPass(&gEntryBlockFinalizerPid)
	{
	}

	virtual const char* getPassName() const {
		return "EntryBlockFinalizer";
	}

	// this is our level 2 semantical analyser on instruction level
    virtual bool runOnFunction(Function &F) {
    	Function::BasicBlockListType& blocks = F.getBasicBlockList();

    	BasicBlock* entry = &blocks.front();
    	BasicBlock* body  = blocks.getNext(entry);

    	Codegen* g = Codegen::get();
    	Builder* builder = g->getBuilder();

    	if (!entry)
    	{
    		cerr << "Function '" << F.getNameStr() << "' doesn't have an entry basic block!" << endl;
    		return false;
    	}
    	else if(entry->empty())
    	{
    		cerr << "Function " << F.getNameStr() << "'s doesn't any instructions in its basic block!" << endl;
    		builder->SetInsertPoint(entry);
    		builder->CreateBr(body);
    		return true;
    	}
    	else if ( isa<BranchInst>(entry->back()) )
    	{
    		return false;
    	}

    	builder->SetInsertPoint(entry);
    	builder->CreateBr(body);

    	// find a basic block that is unreachable
    	// remove the jump instructions jumping to the block
    	BasicBlock* bb = NULL;
    	for (Function::iterator basicBlock = F.begin(); basicBlock != F.end(); basicBlock++)
		{
			bb = &*basicBlock;

    		if(!basicBlock->getInstList().empty())
    		{
				Instruction* instruction = &basicBlock->getInstList().front();
				if (!isa<UnreachableInst>(instruction))
					continue;

				// remove all breaks to this block!
				for (Function::iterator i = F.begin(); i != F.end(); i++)
					for (BasicBlock::iterator j = i->begin(); j != i->end(); j++)
					{
						if(BranchInst* instruction = cast_or_null<BranchInst>(&*j))
							if (!instruction->isConditional())
							{
								// unconditional branch instruction
								if(cast<BasicBlock>(instruction->getOperand(0)) == bb)
								{
									j->removeFromParent();
									j = i->begin();
								}
							}
							else if(cast<BasicBlock>(j->getOperand(2)) == bb ||
									cast<BasicBlock>(j->getOperand(1)) == bb)
							{
								// set the builder to the conditional branch instruction
								BasicBlock::iterator k = j;

								builder->SetInsertPoint(&*i, --k);
								builder->CreateBr(cast<BasicBlock>(instruction->getOperand(2)));
								j->removeFromParent();
								j = i->begin();
							}
					}
    		}

        	bb->removeFromParent();
        	basicBlock = F.begin();
		}

    	return true;
    }
};

static const int gMAliceFinalizerPid = 0;
class MAliceFinalizer : public llvm::FunctionPass
{
public:
	MAliceFinalizer()
	 : FunctionPass(&gMAliceFinalizerPid)
	{
	}

	virtual const char* getPassName() const {
		return "MAliceFinalizer";
	}


	// this is our level 2 semantical analyser on instruction level
    virtual bool runOnFunction(Function &F) {
    	bool changed = false;

    	Function::BasicBlockListType& blocks = F.getBasicBlockList();
		bool isRet, wasRet = false;
    	for (Function::iterator i = F.begin(); i != F.end(); i++)
    	{
    		for (BasicBlock::iterator j = i->begin(); j != i->end(); j++)
    		{
    			isRet = isa<ReturnInst>(*j);
    			if (wasRet && isRet) {
    				changed = true;
    				j->eraseFromParent();
    				j = i->begin();
    			}
    			wasRet = isRet;
    		}
    	}

    	return changed;
    }
};





/* codegen class */
Codegen::Codegen() :
	mContext(), mVars(), mEntry(NULL)
{
	gCodegen = this;
	mModule = new Module("MAlice", mContext);
	mBuilder = new Builder(mContext);

	// set target machine instructions
	getModule()->setDataLayout("e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64");
	getModule()->setTargetTriple("x86_64-linux-gnu");

	// initialise the data types
	getModule()->addTypeName("number", Type::getInt32Ty(getContext()));
	getModule()->addTypeName("letter", Type::getInt8Ty(getContext()));
	getModule()->addTypeName("sentence", Type::getInt8PtrTy(getContext()));

	// initialise the CRT functions
	{ // printf and scanf
		std::vector<const Type*> args;
		args.push_back(Type::getInt8PtrTy(getContext()));

		FunctionType *FT = FunctionType::get(Type::getInt32Ty(getContext()), args, true);

		//create an external function with no body
		Function *printf = Function::Create(FT, GlobalValue::ExternalLinkage, "printf", getModule());
		printf->setCallingConv(CallingConv::C);

		Function *scanf = Function::Create(FT, GlobalValue::ExternalLinkage, "scanf", getModule());
		scanf ->setCallingConv(CallingConv::C);

		args.clear();
		FT = FunctionType::get(Type::getVoidTy(getContext()), args, false);
		Function *abort = Function::Create(FT, GlobalValue::ExternalLinkage, "abort", getModule());
		abort ->setCallingConv(CallingConv::C);
	}

	{// initialise the constant global variables
		GlobalVariable* fmt;

		fmt = new GlobalVariable(/*Module=*/*getModule(),
		/*Type=*/ArrayType::get(Type::getInt8Ty(getContext()), 3),
		/*isConstant=*/true,
		/*Linkage=*/GlobalValue::PrivateLinkage,
				/*Initializer=*/reinterpret_cast<Constant*> (StringLiteralAst("%d").generate()),
				/*Name=*/".fmt_d");
		fmt->setAlignment(1);
		fmt = new GlobalVariable(/*MocreateSimpliVariableAstdule=*/*getModule(),
		/*Type=*/ArrayType::get(Type::getInt8Ty(getContext()), 53),
		/*isConstant=*/true,
		/*Linkage=*/GlobalValue::PrivateLinkage,
				/*Initializer=*/reinterpret_cast<Constant*> (StringLiteralAst("IOError: Could not read integer from standard input!").generate()),
				/*Name=*/".io.error");
		fmt->setAlignment(1);

		fmt = new GlobalVariable(/*Module=*/*getModule(),
		/*Type=*/ArrayType::get(Type::getInt8Ty(getContext()), 3),
		/*isConstant=*/true,
		/*Linkage=*/GlobalValue::PrivateLinkage,
				/*Initializer=*/reinterpret_cast<Constant*> (StringLiteralAst("%s").generate()),
				/*Name=*/".fmt_s");
		fmt->setAlignment(1);

		fmt = new GlobalVariable(/*Module=*/*getModule(),
		/*Type=*/ArrayType::get(Type::getInt8Ty(getContext()), 5),
		/*isConstant=*/true,
		/*Linkage=*/GlobalValue::PrivateLinkage,
				/*Initializer=*/reinterpret_cast<Constant*> (StringLiteralAst("'%c'").generate()),
				/*Name=*/".fmt_c");
		fmt->setAlignment(1);
	}
}

Codegen::~Codegen()
{
	delete mModule;
	delete mBuilder;
}

Codegen* Codegen::get()
{
	assert (gCodegen != NULL && "No code-generator instance is ever instantiated");
	return gCodegen;
}

LLVMContext& Codegen::getContext()
{
	return mContext;
}

static void getFPM(FunctionPassManager* pm, Codegen::Optimisation optimisation)
{
	int opt = static_cast<int> (optimisation);

	if (opt > 1)
	{
		pm->add(createPromoteMemoryToRegisterPass());
		pm->add(createInstructionCombiningPass());
		pm->add(createConstantPropagationPass());
		pm->add(createCFGSimplificationPass());

		pm->add(createDeadInstEliminationPass());
		pm->add(createDeadCodeEliminationPass());
		pm->add(createDeadStoreEliminationPass());
	}

	if (opt > 1)
	{
		pm->add(createReassociatePass());

		pm->add(createBreakCriticalEdgesPass());
		pm->add(createTailCallEliminationPass());
		pm->add(createTailDuplicationPass());

		pm->add(createLCSSAPass());
		pm->add(createLoopDeletionPass());
		pm->add(createLoopSimplifyPass());
		pm->add(createLoopStrengthReducePass());
		pm->add(createLoopUnswitchPass());
		pm->add(createLoopUnrollPass());
		pm->add(createLoopRotatePass());
		pm->add(createLoopIndexSplitPass());
	}

	if (opt > 2)
	{
		// pm->add(new TargetData(getModule()));
		pm->add(createSSIPass());
		pm->add(createSSIEverythingPass());
		pm->add(createBlockPlacementPass());
		pm->add(createSimplifyLibCallsPass());
		pm->add(createSimplifyHalfPowrLibCallsPass());
		pm->add(createSimplifyLibCallsPass());
		pm->add(createScalarReplAggregatesPass());
		pm->add(createGVNPass());
		pm->add(createAggressiveDCEPass());
		pm->add(createGEPSplitterPass());
		pm->add(createSCCVNPass());
		pm->add(createABCDPass());
		pm->add(createSCCPPass());
		pm->add(createLICMPass());

		pm->add(createMemCpyOptPass());

		pm->add(createLowerSwitchPass());
		pm->add(createLowerInvokePass());
	}
}

Module* Codegen::finalise(MAliceDriver& driver, Optimisation optimisation)
{
	Module* module = getModule();
	auto_ptr<FunctionPassManager> pm = auto_ptr<FunctionPassManager>(
			new FunctionPassManager(module));

	// generate the module
	driver.get()->generate();

	// always include instruction namer pass
	pm->add(createInstructionNamerPass());

	// run our own new optimisation passes
	pm->add(new EntryBlockFinalizer());

	// add optimisation specific optimisations
	getFPM(pm.get(), optimisation);

	// remove duplicate returns
	pm->add(new MAliceFinalizer());

	// perform optimisation
	cout << "\tEntering optimisation level " << optimisation << endl;
	for (Module::iterator i = module->begin(); i != module->end(); i++)
		pm->run(*i);

	return module;
}

Builder* Codegen::getBuilder()
{
	return mBuilder;
}

Module* Codegen::getModule()
{
	return mModule;
}


string Codegen::getScopeError(const string& name)
{
	stringstream ss;

	int size = mVars.size();
	ss << "Variable '" << name << "' not found on this scope: " << endl;
	for (vector<ScopeType>::reverse_iterator j = mVars.rbegin(); j != mVars.rend(); j++, size--)
	{
		ss << setw(4 * size) << "{" << endl;
		for (ScopeType::iterator i = j->begin(); i != j->end(); i++)
		{
			ss << setw(4 * size) << i->first << ", " << endl;
		}
		ss << setw(4 * size) << "}" << endl;
	}

	return ss.str();
}

ExpressionAst* Codegen::getDefaultValue(const string& type)
{

	// default value
	if (type == "number")
	{
		return new IntegerLiteralAst(0);
	}
	else if (type == "letter")
	{
		return new CharacterLiteralAst('\0');
	}
	else if (type == "sentence")
	{
		return new StringLiteralAst("");
	}
	else
	{
		cerr << "Warning: could not deduce the type of variable " << type << ". Expect subsequent failures!"<< endl;
		return NULL;
	}
}

Value* Codegen::getVar(const string& name)
{
	for (vector<ScopeType>::reverse_iterator i = mVars.rbegin(); i != mVars.rend(); i++)
	{
		if (Value* value = (*i) [name])
			return value;
	}

	throw getScopeError(name);
}

void Codegen::pushVar(const string& name, Value* value)
{
	ScopeType& top = mVars.back();

	// ensure we are not redefining
	if (top[name] && top[name] != value)
	{
		throw runtime_error(name + " cannot be redefined!\n" + getScopeError(name));
	}

	top[name] = value;
}

void Codegen::pushScope()
{
	mVars.push_back(ScopeType());
}

void Codegen::popScope()
{
	// TODO the scope is exiting, we may want to add destructors
	// and invoke them at scope exit
	mVars.pop_back();
}

}
